// SPDX-License-Identifier: GPL-2.0+

#include <common.h>
#include <command.h>

enum update_dev {
	UPDATE_DEV_UNKNOWN = 0,
	UPDATE_DEV_USB,
	UPDATE_DEV_TFTP,
	UPDATE_DEV_MMC,

	UPDATE_DEV_COUNT,			/* Number of dev types */
};

static const char *update_devname_str[UPDATE_DEV_COUNT] = {
	[UPDATE_DEV_USB]	= "usb",
	[UPDATE_DEV_TFTP]	= "tftp",
	[UPDATE_DEV_MMC]	= "mmc",
};

/* Update types: */
enum update_type {
	UPDATE_TYPE_UNKNOWN = 0,
	UPDATE_TYPE_KERNEL,
	UPDATE_TYPE_ROOTFS,
	UPDATE_TYPE_UBOOT,
	UPDATE_TYPE_RTTHREAD,
	UPDATE_TYPE_ALL,

	UPDATE_TYPE_COUNT,			/* Number of update types */
};

static const char *update_typename_str[UPDATE_TYPE_COUNT] = {
	[UPDATE_TYPE_KERNEL]	= "kernel",
	[UPDATE_TYPE_ROOTFS]	= "rootfs",
	[UPDATE_TYPE_UBOOT]		= "uboot",
	[UPDATE_TYPE_RTTHREAD]		= "rtthread",
	[UPDATE_TYPE_ALL]		= "all",
};

static int update_uboot(int dev)
{
	int ret = -1;
	char addr[80] = {0};

	printf("update u-boot.............\n");
	sprintf(addr, "%lx", 0x80100000);
	setenv("loadaddr", addr);

	switch (dev) {
	case UPDATE_DEV_USB:
		ret = run_command("fatload usb 0:1 ${loadaddr} /update/u-boot.bin", 0);
		break;
	case UPDATE_DEV_TFTP:
		sprintf(addr, "%lx", 0xa0100000);
		setenv("loadaddr", addr);
		ret = run_command("tftpboot ${loadaddr} u-boot.bin", 0);
		sprintf(addr, "%lx", 0x80100000);
		setenv("loadaddr", addr);
		break;
	case UPDATE_DEV_MMC:
		ret = run_command("fatload mmc 0:1 ${loadaddr} /update/u-boot.bin", 0);
		break;
	}

	if (ret) {
		printf("upgrade uboot failed no update/u-boot.bin file!\r\n");
		return ret;
	}

	ret = run_command("sf probe;sf erase 0 0x80000;sf update ${loadaddr} 0 ${filesize}", 0);
	if (ret) {
		printf("upgrade uboot failed!\r\n");
	} else {
		printf("upgrade uboot success\r\n");
	}

	return ret;
}

static int update_kernel(int dev, char *typename)
{
	int ret = -1;
	char addr[80] = {0};

	printf("update kernel....................\n");
	sprintf(addr, "%lx", 0x80100000);
	setenv("loadaddr", addr);

	switch (dev) {
	case UPDATE_DEV_USB:
		ret = run_command("fatload usb 0:1 ${loadaddr} /update/uImage", 0);
		break;
	case UPDATE_DEV_TFTP:
		sprintf(addr, "%lx", 0xa0100000);
		setenv("loadaddr", addr);
		ret = run_command("tftpboot ${loadaddr} uImage", 0);
		sprintf(addr, "%lx", 0x80100000);
		setenv("loadaddr", addr);
		break;
	case UPDATE_DEV_MMC:
		ret = run_command("fatload mmc 0:1 ${loadaddr} /update/uImage", 0);
		break;
	}

	if (ret) {
		printf("upgrade kernel failed no /update/uImage file!\r\n");
		return ret;
	}

	if (strcmp(typename, "nand") == 0) {
		printf("update kernel to nand............\n");
		ret = run_command("sf probe;nand erase.part kernel;nand write ${loadaddr} kernel ${filesize}", 0);
		if (ret) {
			goto err;
		}
	} else if (strcmp(typename, "sata") == 0) {
		printf("update kernel to ssd.............\n");
		ret = run_command("scsi reset;ext4load scsi 0:2 ${loadaddr} uImage", 0);
		if (ret) {
			goto err;
		}
		ret = run_command("scsi reset;ext4write scsi 0:2 ${loadaddr} /uImage ${filesize}", 0);
		if (ret) {
			goto err;
		}
	} else {
		printf("upgrade kernel failed: nand or sata?\r\n");
		return -1;
	}

	printf("upgrade kernel success\r\n");

	return ret;
err:
	printf("upgrade kernel failed!\r\n");
	return ret;
}

static int update_rtthread(int dev, char *typename)
{
	int ret = -1;
	char addr[80] = {0};

	printf("update rt-thread....................\n");
	sprintf(addr, "%lx", 0x80100000);
	setenv("loadaddr", addr);

	switch (dev) {
	case UPDATE_DEV_USB:
		ret = run_command("fatload usb 0:1 ${loadaddr} /update/rtthread.elf", 0);
		break;
	case UPDATE_DEV_TFTP:
		sprintf(addr, "%lx", 0xa0100000);
		setenv("loadaddr", addr);
		ret = run_command("tftpboot ${loadaddr} rtthread.elf", 0);
		sprintf(addr, "%lx", 0x80100000);
		setenv("loadaddr", addr);
		break;
	case UPDATE_DEV_MMC:
		ret = run_command("fatload mmc 0:1 ${loadaddr} /update/rtthread.elf", 0);
		break;
	}

	if (ret) {
		printf("upgrade rt-thread failed no /update/rtthread.elf file!\r\n");
		return ret;
	}

	if (strcmp(typename, "nand") == 0) {
		printf("update kernel to nand............\n");
		ret = run_command("sf probe;nand erase.part kernel;nand write ${loadaddr} kernel ${filesize}", 0);
		if (ret) {
			goto err;
		}
	} else {
		printf("upgrade rt-thread failed: nand\r\n");
		return -1;
	}

	printf("upgrade rt-thread success\r\n");

	return ret;
err:
	printf("upgrade rt-thread failed!\r\n");
	return ret;
}

static int update_rootfs(int dev, char *typename)
{
	int ret = -1;
	char addr[80] = {0};

	printf("update rootfs to nand............\n");
	sprintf(addr, "%lx", 0x80100000);
	setenv("loadaddr", addr);

	switch (dev) {
	case UPDATE_DEV_USB:
#ifdef CONFIG_LS1B_ACCESS_BOARD
		ret = run_command("fatload usb 0:1 ${loadaddr} /update/rootfs-yaffs2.img", 0);
#else
		ret = run_command("fatload usb 0:1 ${loadaddr} /update/rootfs-ubifs-ze.img", 0);
#endif
		break;
	case UPDATE_DEV_TFTP:
		sprintf(addr, "%lx", 0xa0100000);
		setenv("loadaddr", addr);
#ifdef CONFIG_LS1B_ACCESS_BOARD
		ret = run_command("tftpboot ${loadaddr} rootfs-yaffs2.img", 0);
#else
		ret = run_command("tftpboot ${loadaddr} rootfs-ubifs-ze.img", 0);
#endif
		sprintf(addr, "%lx", 0x80100000);
		setenv("loadaddr", addr);
		break;
	case UPDATE_DEV_MMC:
#ifdef CONFIG_LS1B_ACCESS_BOARD
		ret = run_command("fatload mmc 0:1 ${loadaddr} /update/rootfs-yaffs2.img", 0);
#else
		ret = run_command("fatload mmc 0:1 ${loadaddr} /update/rootfs-ubifs-ze.img", 0);
#endif
		break;
	}

	if (ret) {
#ifdef CONFIG_LS1B_ACCESS_BOARD
		printf("upgrade rootfs failed no /update/rootfs-yaffs2.img file!\r\n");
#else
		printf("upgrade rootfs failed no /update/rootfs-ubifs-ze.img file!\r\n");
#endif
		return ret;
	}

	ret = run_command("sf probe;nand erase.part root;nand write ${loadaddr} root ${filesize}", 0);
	if (ret) {
		printf("upgrade rootfs failed!\r\n");
	} else {
		printf("upgrade rootfs success\r\n");
	}

	return ret;
}

static int do_loongson_update(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
{
	int ret;
	int dev = UPDATE_DEV_UNKNOWN;
	int type = UPDATE_TYPE_UNKNOWN;
	char addr[80] = {0};

	for (dev = 0; dev < UPDATE_DEV_COUNT; dev++) {
		if (update_devname_str[dev] &&
			!strcmp(argv[1], update_devname_str[dev]))
			break;
	}
	if ((dev == UPDATE_DEV_UNKNOWN) || (dev >= UPDATE_DEV_COUNT)) {
		printf("upgrade dev unknown\r\n");
		printf("please intput dev: usb\\tftp\\mmc\r\n");
		goto err;
	}

	for (type = 0; type < UPDATE_TYPE_COUNT; type++) {
		if (update_typename_str[type] &&
			!strcmp(argv[2], update_typename_str[type]))
			break;
	}
	if ((type == UPDATE_TYPE_UNKNOWN) || (type >= UPDATE_TYPE_COUNT)) {
		printf("upgrade type unknown\r\n");
		printf("please intput type: kernel\\rootfs\\uboot\\all\\rt-thread\r\n");
		goto err;
	}

	switch (dev) {
	case UPDATE_DEV_USB:
		run_command("usb reset", 0);
		ret = run_command("usb storage", 0);
		if (ret) {
			printf("upgrade failed no usb storage!\r\n");
			goto err;
		}
		break;
	case UPDATE_DEV_TFTP:
		break;
	case UPDATE_DEV_MMC:
		ret = run_command("mmc rescan", 0);
		if (ret) {
			printf("upgrade failed no mmc storage!\r\n");
			goto err;
		}
		break;
	default:
		break;
	}

	switch (type) {
	case UPDATE_TYPE_UBOOT:
		ret = update_uboot(dev);
		break;
	case UPDATE_TYPE_KERNEL:
		ret = update_kernel(dev, argv[3]);
		break;
	case UPDATE_TYPE_ROOTFS:
		ret = update_rootfs(dev, NULL);
		break;
	case UPDATE_TYPE_ALL:
		ret = update_uboot(dev);
		ret = update_kernel(dev, argv[3]);
		ret = update_rootfs(dev, NULL);
		break;
	case UPDATE_TYPE_RTTHREAD:
		ret = update_rtthread(dev, argv[3]);
		break;
	default:
		break;
	}

	sprintf(addr, "%lx", CONFIG_SYS_LOAD_ADDR);
	setenv("loadaddr", addr);

	return ret;

err:
	return -1;
}

U_BOOT_CMD(
	loongson_update,    4,    1,     do_loongson_update,
	"upgrade uboot kernel rootfs rt-thread over usb\\tftp\\mmc.",
	""
);
